package v1

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"
)

// 分组状态
const (
	GroupStateIdle      = 0
	GroupStateHasTask   = 1
	GroupStateExecuting = 2
)

// ========================================
// Group：单个分组
// ========================================
type Group struct {
	chain       *TaskChain
	state       int32      // 原子
	mu          sync.Mutex // 保护条件变量
	cond        sync.Cond
	groupID     interface{}
	parent      *GroupTask
	maxCount    int
	owner       int64       // 原子
	needCleanup atomic.Bool // 标记是否需要清理

	operaNum atomic.Int32 //
}

func NewGroup(groupID interface{}, parent *GroupTask) *Group {
	g := &Group{
		chain:   NewTaskChain(),
		state:   GroupStateIdle,
		groupID: groupID,
		parent:  parent,
	}
	g.chain.maxCount = 1024
	g.cond.L = &g.mu
	return g
}

// ---- 状态封装 ----
func (g *Group) getState() int32   { return atomic.LoadInt32(&g.state) }
func (g *Group) setState(s int32)  { atomic.StoreInt32(&g.state, s) }
func (g *Group) getOwner() int64   { return atomic.LoadInt64(&g.owner) }
func (g *Group) setOwner(id int64) { atomic.StoreInt64(&g.owner, id) }
func (g *Group) casOwner(old, new int64) bool {
	return atomic.CompareAndSwapInt64(&g.owner, old, new)
}

// --------------------------------------
func (g *Group) AddTask(cb func(args ...interface{}), args ...interface{}) error {

	// 确保没有被删除
	for {
		n := g.operaNum.Load()
		if n < 0 {
			return fmt.Errorf("group is invalid")
		}
		if g.operaNum.CompareAndSwap(n, n+1) {
			break
		}
	}
	defer g.operaNum.Add(-1)

	err := g.chain.Add(cb, args...)
	if err != nil {
		return err
	}
	g.setState(GroupStateHasTask)
	g.cond.Signal() // 唤醒本 group 的 worker
	return nil
}

// 尝试获取处理权
func (g *Group) acquire(workerID int64) bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	if (g.getState() == GroupStateIdle || g.getState() == GroupStateHasTask) &&
		g.casOwner(0, workerID) {
		return true
	}
	return false
}

func (g *Group) release() { g.setOwner(0) }

// 主循环
func (g *Group) processTasks(workerID int64) {
	defer g.release()
	for g.parent.isRunning() && g.operaNum.Load() != -1 { // g.operaNum.Load() 表示分组被清理了
		if g.chain.IsEmpty() {
			g.mu.Lock()
			g.setState(GroupStateIdle)
			// 可能被 Stop 广播唤醒，也可能被 AddTask 信号唤醒
			g.cond.Wait()
			g.mu.Unlock()
			continue
		}
		g.setState(GroupStateExecuting)
		t := g.chain.Next()
		if t != nil {
			func() {
				defer func() {
					g.parent.incCompleted()
				}()
				t.cb(t.args...)
			}()
		}
		// 任务链仍有任务则保持 HasTask，否则 Idle 并通知父对象
		if g.chain.IsEmpty() {
			g.setState(GroupStateIdle)
			g.needCleanup.Store(true) // 再次标记需要清理
			g.parent.releaseWorker(g.groupID)
			break
		} else {
			g.setState(GroupStateHasTask)
		}
	}
}

// ========================================
// GroupTask：调度器
// ========================================
type GroupTask struct {
	currentWorkers    atomic.Int32
	busyWorkers       atomic.Int32
	minWorkers        int
	maxWorkers        int
	groups            sync.Map
	workerCount       sync.WaitGroup
	running           int32 // 原子，1 运行 0 停止
	totalTasks        int64
	completedTasks    int64
	taskNotify        chan struct{}
	cond              sync.Cond // 用于 WaitAllTasks
	workerGroups      sync.Map  // workerID -> groupID
	workerIDSeq       int64
	groupDeleteTimes  atomic.Int32
	maxGroupQueueSize int32
}

func NewGroupTask(minWorkers, maxWorkers int) (*GroupTask, error) {
	if minWorkers <= 0 || maxWorkers <= 0 || minWorkers > maxWorkers {
		return nil, errors.New("invalid worker count parameters")
	}
	gt := &GroupTask{
		minWorkers:        minWorkers,
		maxWorkers:        maxWorkers,
		running:           1,
		taskNotify:        make(chan struct{}, maxWorkers),
		maxGroupQueueSize: 8192,
	}
	gt.cond.L = &sync.Mutex{}
	for i := 0; i < minWorkers; i++ {
		gt.startWorker()
	}
	return gt, nil
}

func (gt *GroupTask) isRunning() bool { return atomic.LoadInt32(&gt.running) == 1 }

func (gt *GroupTask) incCompleted() {
	atomic.AddInt64(&gt.completedTasks, 1)
	gt.cond.Broadcast()
}

func (gt *GroupTask) GetGroupCount() int {
	n := 0
	gt.groups.Range(func(key, value interface{}) bool {
		n++
		return true
	})
	return n
}

// 启动一个 worker
func (gt *GroupTask) startWorker() {
	gt.workerCount.Add(1)
	gt.currentWorkers.Add(1)
	id := atomic.AddInt64(&gt.workerIDSeq, 1)
	go func(workerID int64) {
		defer func() {
			gt.currentWorkers.Add(-1)
			gt.workerGroups.Delete(workerID)
			gt.workerCount.Done()
		}()
		for gt.isRunning() {
			gid, g := gt.getGroupToProcess(workerID)
			if g != nil {
				gt.busyWorkers.Add(1)
				gt.workerGroups.Store(workerID, gid)
				g.processTasks(workerID)
				gt.workerGroups.Delete(workerID)
				gt.busyWorkers.Add(-1)
			} else {
				select {
				case <-gt.taskNotify:
				case <-time.After(100 * time.Millisecond):
					if gt.currentWorkers.Load() > int32(gt.minWorkers) {
						return
					}
				}
			}
		}
	}(id)
}

// 取一个待处理分组
func (gt *GroupTask) getGroupToProcess(workerID int64) (interface{}, *Group) {
	var foundGID interface{}
	var foundG *Group
	gt.groups.Range(func(k, v interface{}) bool {
		g := v.(*Group)
		if g.getState() == GroupStateHasTask && g.acquire(workerID) {
			foundGID, foundG = k, g
			return false
		}
		return true
	})
	return foundGID, foundG
}

// 分组任务执行完，尝试回收或补充线程
func (gt *GroupTask) releaseWorker(groupID interface{}) {
	// 快速看一遍是否还有 HasTask
	has := false
	gt.groups.Range(func(k, v interface{}) bool {
		if v.(*Group).getState() == GroupStateHasTask {
			has = true
			return false
		}
		return true
	})
	if has && gt.currentWorkers.Load() < int32(gt.maxWorkers) {
		// 再看是否有真正空闲的 worker（未绑定 group）
		idle := false
		gt.workerGroups.Range(func(k, v interface{}) bool {
			if v == nil {
				idle = true
				return false
			}
			return true
		})
		if !idle {
			gt.startWorker()
		}
	}
}

// --------------------------------------
// 任务投递
// --------------------------------------
func (gt *GroupTask) PostTaskArgsEx(groupID interface{},
	cb func(args ...interface{}), args ...interface{}) error {
	if !gt.isRunning() {
		return errors.New("GroupTask is stopped")
	}
	if cb == nil {
		return errors.New("callback function cannot be nil")
	}
	value, _ := gt.groups.LoadOrStore(groupID, NewGroup(groupID, gt))
	g := value.(*Group)
	err := g.AddTask(cb, args...)
	if err != nil {
		return err
	}
	atomic.AddInt64(&gt.totalTasks, 1)

	// 需要更多线程？
	if gt.currentWorkers.Load() < int32(gt.maxWorkers) {
		idle := false
		idle = gt.currentWorkers.Load() > gt.busyWorkers.Load()
		//gt.workerGroups.Range(func(_, v interface{}) bool {
		//	if v == nil {
		//		idle = true
		//		return false
		//	}
		//	return true
		//})
		if !idle {
			gt.startWorker()
		}
	}
	select {
	case gt.taskNotify <- struct{}{}:
	default:
	}
	return nil
}

func (gt *GroupTask) PostTaskFunc(groupID interface{}, cb func()) error {
	if cb == nil {
		return errors.New("callback function cannot be nil")
	}
	return gt.PostTaskArgsEx(groupID, func(args ...interface{}) { cb() })
}

// --------------------------------------
// 停止与等待
// --------------------------------------
func (gt *GroupTask) Stop() {
	if !atomic.CompareAndSwapInt32(&gt.running, 1, 0) {
		return // 已停
	}
	close(gt.taskNotify)
	// 广播所有 group.cond
	gt.groups.Range(func(_, v interface{}) bool {
		v.(*Group).cond.Broadcast()
		return true
	})
	// 广播 WaitAllTasks
	gt.cond.Broadcast()
	gt.workerCount.Wait()
}

func (gt *GroupTask) GetCurrentWorkNum() int {
	return int(gt.currentWorkers.Load())
}

// WaitAllTasks 等待所有任务完成；如已 Stop 也会立即返回
func (gt *GroupTask) WaitAllTasks() {
	gt.cond.L.Lock()
	defer gt.cond.L.Unlock()
	for atomic.LoadInt64(&gt.completedTasks) < atomic.LoadInt64(&gt.totalTasks) && gt.isRunning() {
		gt.cond.Wait()
	}
}

// 可选：增加定时清理机制，避免长期未处理的空分组
func (gt *GroupTask) StartCleanupTimer(ctx context.Context, interval time.Duration) {

	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for gt.isRunning() {
		select {
		case <-ticker.C:
			gt.cleanupEmptyGroups()
		case <-ctx.Done():
			return
		}
	}

}

// 在 GroupTask 中增加清理空分组的方法
func (gt *GroupTask) cleanupEmptyGroups() {
	gt.groups.Range(func(k, v interface{}) bool {
		g := v.(*Group)
		// 检查分组是否空闲且标记需要清理
		if g.getState() == GroupStateIdle && g.needCleanup.Load() {
			// 双重检查：确保清理前没有新任务加入
			g.mu.Lock()
			df := false
			if g.operaNum.CompareAndSwap(0, -1) {
				if g.chain.IsEmpty() {
					g.cond.Broadcast()
					gt.groups.Delete(k)
					gt.groupDeleteTimes.Add(1)
					df = true
				} else {
					g.operaNum.Store(0)
				}
			}
			if !df {
				g.needCleanup.Store(false) // 有新任务，取消清理标记
			}
			g.mu.Unlock()
		}
		return true
	})
}

// WaitAllTasksWithCtx 新增：支持外部超时
func (gt *GroupTask) WaitAllTasksWithCtx(ctx context.Context) error {
	gt.cond.L.Lock()
	defer gt.cond.L.Unlock()
	for atomic.LoadInt64(&gt.completedTasks) < atomic.LoadInt64(&gt.totalTasks) && gt.isRunning() {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		gt.cond.Wait()
	}
	return nil
}

var (
	DefaultGroupTask, _ = NewGroupTask(1, 5)
)
